library(ggplot2)
library(gridExtra)

## plot rho_t

rho_plot <- function(out,  rho, constant = FALSE) {
	
	modelRho <- datap <- NULL

	if (constant == FALSE) {
		modelRho <- out$Rho
		TT <- dim(modelRho)[1]
		## ----------------------- 1 varying rho_t
		datap <- cbind.data.frame(1:TT, apply(modelRho, 1, mean),
		                          apply(modelRho, 1, quantile, 0.025),
		                          apply(modelRho, 1, quantile, 0.975),
		                          c(rho))
	} else {
		modelRho <- out$rho0
		datap <- cbind.data.frame(1:TT, rep(mean(modelRho), TT),
		                         rep(quantile(modelRho, 0.025), TT),
		                         rep(quantile(modelRho, 0.975), TT),
		                         c(rho))
	}
	

	names(datap) <- c("time", "mean", "cil", "ciu", "true")

	datap$type <- rep("m", TT)

	p <- ggplot(datap) + xlab("Time") + ylab("Rho")

	p <- p + geom_line(aes(time, mean, colour = type, size = type, linetype = type))
	p <- p + geom_line(aes(time, true), size = 1, colour = "#000000")
	p <- p + geom_ribbon(aes(x = time, ymin = cil, 
	                                   ymax = ciu), fill = "#000000", alpha = 0.2)

	set.limits <- c("t", "m", "c")
	set.labels <- c("True", "Posterior Mean", "95% CI")
	set.colors <- c("#000000", "grey50", "#00000080")
	set.size <- c(1, 1, 3)
	set.type <- c("solid", "dashed", "solid")

	p <- p + scale_colour_manual(limits = set.limits,
	                             labels = set.labels,
	                             values =set.colors) +
	         scale_size_manual(limits = set.limits,
	                           labels = set.labels,
	                           values = set.size) +
	         scale_linetype_manual(limits = set.limits,
	                           labels = set.labels,
	                           values = set.type) +
	         guides(colour = guide_legend(title=NULL, nrow=1),
	                size = guide_legend(title=NULL, nrow=1),
	                linetype = guide_legend(title=NULL, nrow=1))


	## p <- p + theme_bw(base_size = 15)
	p <- p + theme(
	               legend.text = element_text(margin = margin(r = 10, unit = "pt"), size = 12),
	               legend.position = "bottom",
	               axis.title = element_text(size=12),
	               axis.title.x = element_text(margin = margin(t = 8, r = 0, b = 0, l = 0)),
	               axis.title.y = element_text(margin = margin(t = 0, r = 8, b = 0, l = 0)),
	               axis.text = element_text(color="black", size=8),
	               axis.text.x = element_text(size = 8),
	               axis.text.y = element_text(size = 8),
	               plot.title = element_text(size = 15,
	                                         hjust = 0.5,
	                                         face="bold",
	                                         margin = margin(10, 0, 10, 0)))


	return(p) 

}

## plot beta
beta_plot <- function(data, tb, pos) {
	## beta1
	subdata <- data[which(data[, "beta"] == pos), ]
	trueb <- tb ## 2 for beta3

	scaleFUN <- function(x) sprintf("%.2f", x) ## integer value at x axis

	p <- ggplot(subdata, aes(x = pos, y = mean, group = model, color = model)) 
	p <- p + xlab("Model") + ylab(paste("Beta", pos, sep = ""))
	p <- p + geom_hline(yintercept = trueb, linetype = "dashed", colour = "grey50")
	p <- p + geom_pointrange(aes(ymin=ci_l, ymax=ci_u))
	p <- p + scale_y_continuous(labels=scaleFUN)


	p <- p + theme(legend.text = element_text(margin = margin(r = 10, unit = "pt"), size = 15),
	               legend.position = "bottom",
	               axis.title = element_text(size=12),
	               axis.title.y = element_text(margin = margin(t = 0, r = 8, b = 0, l = 0)),
	               axis.text = element_text(color="black", size=8),
	               axis.text.x=element_blank(),
	               axis.ticks.x=element_blank(),
	               axis.title.x=element_blank(),
	               axis.text.y = element_text(size = 8),
	               plot.title = element_text(size = 15,
	                                         hjust = 0.5,
	                                         face="bold",
	                                         margin = margin(10, 0, 10, 0)))


	return(p)
}

## plot omega 
omegaPlot <- function(x,                      ## summary data
                      oxlim = NULL,
                      plotlength = 5,         ## for multi-level and time-varying, by row or column
                      labelname = NULL){      ## annotation




    data <- x$omega
    
    k <- dim(data)[1]
    p.all <- NULL

    for (i in 1:k) {
      	subdata <- cbind.data.frame(data[i,])
      	names(subdata) <- "omega"

      	p <- ggplot(subdata,aes(omega))
        p <- p + geom_vline(xintercept = 0, linetype = "dashed", colour = "grey50")

        p <- p + geom_line(stat = "density") + expand_limits(y = 0)
        if (!is.null(oxlim)) {
            p <- p + xlim(oxlim[1], oxlim[2]) 
        }

      	p <- p + xlab("Omega") + ylab("Density")
      	p <- p + theme_bw(base_size = 15)
      	p <- p + theme(panel.grid.major = element_blank(),
      		             panel.grid.minor = element_blank(),
      	            	 legend.text = element_text(margin = margin(r = 10, unit = "pt"), size = 12),
      	               legend.position = "bottom",
                       axis.title = element_text(size=12),
                       axis.title.x = element_text(margin = margin(t = 8, r = 0, b = 0, l = 0)),
                       axis.title.y = element_text(margin = margin(t = 0, r = 8, b = 0, l = 0)),
                       axis.text = element_text(color="black", size=8),
                       axis.text.x = element_text(size = 8),
                       axis.text.y = element_text(size = 8),
                       plot.title = element_text(size = 15,
                                                 hjust = 0.5,
                                                 face="bold",
                                                 margin = margin(10, 0, 10, 0)))
        
        
        main <- ifelse(is.null(labelname), "", paste(labelname, i, sep = "")) 
        p <- p + ggtitle(main)
      	p.all <- c(p.all, list(p))

    }

    q <- grid.arrange(grobs = p.all, nrow = plotlength)
}
